run_name: megablocks-matmul-benchmark
image: mosaicml/llm-foundry:2.0.1_cu118-latest
cluster: r9z1
gpu_num: 8
gpu_type: h100_80gb
integrations:
  - integration_type: git_repo
    git_repo: stanford-futuredata/megablocks
    git_branch: main
    pip_install: absl-py 'git+https://github.com/openai/triton.git@main#egg=triton&subdirectory=python'
    ssh_clone: false
command: |-
  cd megablocks
  export ENABLE_TMA=1
  export ENABLE_MMA_V3=1

  pip install git+https://github.com/stanford-futuredata/stk.git@main --no-dependencies
  pip install git+https://github.com/tgale96/grouped_gemm.git@main --no-dependencies
  pip install . --no-dependencies

  python megablocks/ops/matmul_benchmark.py
